export CUDA_LAUNCH_BLOCKING=1
export NCCL_DEBUG=INFO

deepspeed --num_gpus 1 main.py \
   --data_path trl-lib/hh-rlhf-helpful-base \
   --data_split 4,5,1 \
   --actor_model_name_or_path /root/autodl-tmp/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step1_supervised_finetuning/qwen3-4b-base-lora-sft-helpful-merged \
   --critic_model_name_or_path /root/autodl-tmp/DeepSpeedExamples/applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/helpful-rm \
   --num_padding_at_beginning 0 \
   --per_device_generation_batch_size 4 \
   --per_device_training_batch_size 4 \
   --generation_batches 1 \
   --ppo_epochs 1 \
   --max_answer_seq_len 256 \
   --max_prompt_seq_len 256 \
   --num_train_epochs 1 \
   --actor_zero_stage 2 \
   --critic_zero_stage 2 \
   --offload \
   --lr_scheduler_type cosine \
   --gradient_accumulation_steps 4 \
   --dtype bf16 \
   --num_warmup_steps 100 \
   --deepspeed --seed 1234 \
   --actor_lora_dim 32 \
   --actor_lora_module_name "q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj, gate_proj, lm_head" \
   --only_optimize_lora \
   --actor_gradient_checkpointing \
   --critic_gradient_checkpointing \
   --enable_hybrid_engine \
   --actor_dropout 0.0 \
   --output_dir ./output \
    &> ./output/training.log

   # --actor_zero_stage 2 \
   # --offload_reference_model \
   # --offload \
      # --dtype 'fp16' \
         # --offload \